#include <iostream>

using namespace std;

#ifdef _MSC_VER
#pragma warning(disable : 4996)
#endif

static size_t read_count = 0;
int read_data(FILE* fp, unsigned char* buff, int len = 1)
{
    size_t n = fread(buff, 1, len, fp);
    read_count += n;
    if (n < len)
        return -1;
    return 0;
}

#define GETDATA(fp, c, len) \
{\
if (read_data(fp, c, len) < 0) {\
    printf("end of file\n");\
    return -1;\
}\
}

typedef struct _EBSP
{
    unsigned char* data; // ebsp data
    int ebsp_len;        // ebsp length

    int complete;                    // private data 

    unsigned char remain_data[1024]; // private data 
    int remain_data_len;             // private data 

    unsigned char start_code[5];     // private data 
    int start_code_len;              // private data 

}EBSP;

typedef struct _RBSP
{
    unsigned char* data;
    int rbsp_len;

}RBSP;

typedef struct _BITSTREAM {

    unsigned char* data;
    int data_len;
    int bit_index;
}BITSTREAM;

typedef struct _SPSNALU
{
    int profile_idc;
    char profile_idc_name[128];
    int width;
    int high;
    int fps;

}SPSNALU;

int find_start_code(unsigned char* _data, EBSP* ebsp, int len)
{
 
    int i = 0;
    while (i < len-2)
    {
        if (_data[i] == 0x00 && _data[i + 1] == 0x00) {
        
            if (_data[i+2] == 0x01)
            {
                ebsp->start_code[0] = 0x00;
                ebsp->start_code[1] = 0x00;
                ebsp->start_code[2] = 0x01;
                ebsp->start_code[3] = _data[i + 3];
                ebsp->start_code_len = 3;

                return i + 2+1;
            }

            if (_data[i + 2] == 0x00 && _data[i + 3] == 0x01)
            {
                ebsp->start_code[0] = 0x00;
                ebsp->start_code[1] = 0x00;
                ebsp->start_code[2] = 0x00;
                ebsp->start_code[3] = 0x01;
                ebsp->start_code[4] = _data[i + 4];
                ebsp->start_code_len = 4;
                return i + 3+1;
            }
        }
        i++;
    }

    return 0;
}

void init_ebsp(EBSP* ebsp)
{
    memset(ebsp, 0x00, sizeof(EBSP));
    ebsp->ebsp_len = 0;
    ebsp->data = (unsigned char*)malloc(1);
}

int get_ebsp(FILE* fp, EBSP* ebsp)
{
    if (ebsp->remain_data_len < 1024) {
        GETDATA(fp, ebsp->remain_data+ ebsp->remain_data_len, 1024 - ebsp->remain_data_len);
        ebsp->remain_data_len = 1024;
    }

    if (ebsp->complete == 1) {
    
        ebsp->ebsp_len = 0;
        ebsp->complete = 0;
    }

    int start_pos = 0, end_pos = 0;

    while (1) {
        start_pos = find_start_code(ebsp->remain_data, ebsp, ebsp->remain_data_len); // 找第一个start code
        if (start_pos == 0) { 

            memcpy(ebsp->remain_data, ebsp->remain_data + ebsp->remain_data_len, 2);
            ebsp->remain_data_len = 2;

            GETDATA(fp, ebsp->remain_data + ebsp->remain_data_len, 1024 - ebsp->remain_data_len);
            ebsp->remain_data_len = 1024;
            continue;
        }
        else {   // find first start code
            ebsp->remain_data_len = 1024 - start_pos;
            memcpy(ebsp->remain_data, ebsp->remain_data + start_pos, ebsp->remain_data_len);
        }
        break;
    }

    while (1)
    {
      
        end_pos = find_start_code(ebsp->remain_data, ebsp, ebsp->remain_data_len);
        if (end_pos == 0) {

            int pos = ebsp->ebsp_len;
            ebsp->ebsp_len += (ebsp->remain_data_len-2);
            ebsp->data = (unsigned char*)realloc(ebsp->data, ebsp->ebsp_len);
            memcpy(ebsp->data+pos, ebsp->remain_data, ebsp->remain_data_len-2);
            
            memcpy(ebsp->remain_data, ebsp->remain_data + ebsp->remain_data_len-2, 2);
            ebsp->remain_data_len = 2;

            GETDATA(fp, ebsp->remain_data+ ebsp->remain_data_len, 1024 - ebsp->remain_data_len);
            ebsp->remain_data_len = 1024;

            continue;
        }
        else {
            
            int pos = ebsp->ebsp_len;
            ebsp->ebsp_len += (end_pos - ebsp->start_code_len);
            ebsp->data = (unsigned char*)realloc(ebsp->data, ebsp->ebsp_len);
            memcpy(ebsp->data + pos, ebsp->remain_data , (end_pos - ebsp->start_code_len));

            memcpy(ebsp->remain_data, ebsp->remain_data + end_pos - ebsp->start_code_len, ebsp->remain_data_len - end_pos + ebsp->start_code_len);
            ebsp->remain_data_len = ebsp->remain_data_len - end_pos + ebsp->start_code_len;
            ebsp->complete = 1;
        }
        break;
    }
   // printf("start:%d, end:%d---data_size:%d\n", start_pos, end_pos, rbsp->rbsp_len);
   
    return 0;
}

int get_nalu_type(unsigned char data)
{
    return data & 0x1f;
}

int get_nalu_nri(unsigned char data)
{
    return data>>5 & 0x3;
}

int read_bit(unsigned char* data, int len, BITSTREAM* bit)
{
    int val = 0;
    int i = 0;
    while (i < len) {

        if (bit->bit_index < 0) {
            data++;
            bit->bit_index = 7;
        }

        int bit_data = ((*data) >> bit->bit_index) & 0x01;
        bit->bit_index--;

        val |= bit_data << len - 1 -i ;
        i++;
    }
    return val;
}

int ue(unsigned char* rbsp, BITSTREAM* bit)
{
    int r = 0;
    int length = 0; // 需要读取的数据长度(bit)

    int byte_index = 0;
    while (1) {
        int data = read_bit(rbsp, 1, bit);
        if (data == 0) {
            length++;
            continue;
        }
        else {
        
            break;
        }
    }

    int val = read_bit(rbsp, length, bit);
    val += (1 << length) - 1;

    return val;
}

int parse_nonI_frame(unsigned char* rbsp, int len)
{
    if (rbsp == NULL || len < 2) {
        printf("invalid data\n");
        return -1;
    }
    rbsp++; // skip nalu header
    
    BITSTREAM bit;
    memset(&bit, 0x00, sizeof(bit));
    bit.data = rbsp;
    bit.bit_index = 7;

    int first_mb_in_slice = ue(bit.data, &bit);
    int slice_type = ue(rbsp, &bit);
    //printf("type:%d\n", slice_type);
    return slice_type;
}



int parse_sps(unsigned char* data, SPSNALU* sps)
{
    BITSTREAM bit;
    memset(&bit, 0x00, sizeof(bit));
    bit.data = data;
    bit.bit_index = 7;

    sps->profile_idc = read_bit(data, 8, &bit);
    if (sps->profile_idc == 66)
        strcpy(sps->profile_idc_name, "Baseline");
    else if (sps->profile_idc == 77)
        strcpy(sps->profile_idc_name, "Main");
    else if (sps->profile_idc == 88)
        strcpy(sps->profile_idc_name, "Extened");
    else if (sps->profile_idc == 100)
        strcpy(sps->profile_idc_name, "High");
    else if (sps->profile_idc == 110)
        strcpy(sps->profile_idc_name, "High 10");
    else if (sps->profile_idc == 122)
        strcpy(sps->profile_idc_name, "High 4:2:2");
    else if (sps->profile_idc == 144)
        strcpy(sps->profile_idc_name, "High 4:4:4");


    return 0;
}

void ebsp2rbsp(unsigned char* data, int len, RBSP* rbsp)
{
    rbsp->data = (unsigned char*)realloc(rbsp->data, len);
    rbsp->rbsp_len = 0;

    int rbsp_index = 0;
    int i = 0;
    for (; i < len; i++)
    {
        if (i < 2) {
            rbsp->data[rbsp->rbsp_len++] = data[i];
            continue;
        }

        if (data[i] == 0x03)
        {
            if (data[i - 2] == 0x00 && data[i - 1] == 0x00)
            {
                if (i+1 < len && (data[i + 1] == 0x00 || data[i + 1] == 0x01 || data[i + 1] == 0x02 || data[i + 1] == 0x03) {
                
                    continue;
                }
                else {
                    rbsp->data[rbsp->rbsp_len++] = data[i];
                }
            }
            else {
                rbsp->data[rbsp->rbsp_len++] = data[i];
            }
        }
        else {
            rbsp->data[rbsp->rbsp_len++] = data[i];
        }
    }
    return;
}

int main(int argc, char** argv)
{
    std::cout <<argv[1]<< endl;
    FILE* fp = fopen(argv[1], "rb");
    FILE* fp_out = fopen("rbsp.data", "wb");
    EBSP ebsp;

    init_ebsp(&ebsp);
    int index = 0;

    RBSP rbsp;
    rbsp.data = (unsigned char*)malloc(1);

    while (1) {

        if (get_ebsp(fp, &ebsp) < 0)
            break;

        ebsp2rbsp(ebsp.data, ebsp.ebsp_len, &rbsp);
        printf("%x %x\n", rbsp.data[0], rbsp.data[1]);
        int type = get_nalu_type(rbsp.data[0]);
        int nri = get_nalu_nri(rbsp.data[0]);

        switch (type)
        {
        case 6:
            printf("SEI nalu size:%d\n", rbsp.rbsp_len);
            break;
        case 7:
            SPSNALU sps;
            memset(&sps, 0x00, sizeof(sps));
            parse_sps(rbsp.data+1, &sps);
            printf("SPS nalu size:%d\n", rbsp.rbsp_len);
            printf("profile:%s\n", sps.profile_idc_name);


            break;
        case 8:
            printf("PPS nalu size:%d\n", rbsp.rbsp_len);
            break;
        case 5:
            printf("I slice size:%d--#%d\n", rbsp.rbsp_len, index++);
            break;
        case 1: {
            int slice_type = parse_nonI_frame(rbsp.data, rbsp.rbsp_len);
            if (slice_type == 5 ||slice_type == 0) {
                printf("P slice nri:%d(%s) ---size:%d--#%d\n", nri, (nri==0)?"非参考帧":"参考帧" ,rbsp.rbsp_len, index++);
            }
            else if (slice_type == 6||slice_type == 1) {
            
                printf("B slice nri:%d(%s) ---size:%d--#%d\n", nri, (nri == 0) ? "非参考帧" : "参考帧", rbsp.rbsp_len, index++);
            }
           
            break;
        }
        default:
            printf("error\n");
            break;
        }
        
        fwrite(rbsp.data, rbsp.rbsp_len, 1, fp_out);
        fflush(fp_out);
        //free(rbsp.data);
    }
    printf("exit\n");
    return 0;
}
